from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.append('../gprep')  # path to code
from spline_func import Spline
from prediction_statistics2 import Prediction_Stats
from plot_functions import visualize_training_data, plot_repulsion, plot_RMSE_vs_hyper
from gp_sparse_linear import GP_sparse_linear

# some plotting tools
# class for comparing the prediction performance
# class for reading in repulsion spline from .skf file

# from gp_sparse_linear import GP_sparse_linear # class for Gaussian process regression
# training data
training_file = 'tmp2.rep'
bond_file = 'bond_dists_tmp2.rep'  # list of bond distances in the training file
# (additional output of get_distances.py)
#visualize_training_data(bond_file, 'Hf-Hf', figsize=None)
#visualize_training_data(bond_file, 'Hf-O', figsize=None)
#visualize_training_data(bond_file, 'O-O', figsize=None)

# histogram of bond distances
# visualize_training_data(bond_file, 'C-C', figsize=None)
# visualize_training_data(bond_file, 'C-O', figsize=None)
# visualize_training_data(bond_file, 'C-H', figsize=None)
# visualize_training_data(bond_file, 'O-O', figsize=None)
# visualize_training_data(bond_file, 'O-H', figsize=None)
# visualize_training_data(bond_file, 'H-H', figsize=None)
# the covariance function (kernel) k expresses our prior assumptions about the fit function.
# the parameters of the kernel are called hyperparameters:

# cutoff: range of the repulsive potential
# sigma : data noise (if delta = 1, then data noise to prior function variance ratio) --> prevents overfitting
# delta: prior function variance
# theta: length scale of covariance (over which length are function values similar to each other?)
# beta: exponential damping (for smoother function in the long range)
# X_pseudo: representative bond lengths

# all hyperparameters, except sigma, and the pseudo inputs are bond type specific
# and are passed in a 'bond dictionary'

sigma = 0.05
beta_global = 3.0
theta_global = 1.0

cutoff_HfHf = 6.0
start_HfHf = 1.1  # smallest bond distance in training file
HfHf_dict = {'X_pseudo': np.linspace(start_HfHf, cutoff_HfHf, 20),
           'beta': beta_global,
           'theta': theta_global,
           'cutoff': cutoff_HfHf}

cutoff_HfO = 6
start_HfO = 0.8
HfO_dict = {'X_pseudo': np.linspace(start_HfO, cutoff_HfO, 20),
           'beta': beta_global,
           'theta': theta_global,
           'cutoff': cutoff_HfO}

cutoff_OO = 6
start_OO = 1.0
OO_dict = {'X_pseudo': np.linspace(start_OO, cutoff_OO, 20),
           'beta': beta_global,
           'theta': theta_global,
           'cutoff': cutoff_OO}


bond_dict = {'Hf-Hf': HfHf_dict, 'Hf-O': HfO_dict, 'O-O': OO_dict}

# kernel: function that expresses covariance between function values
# Squared Exponential covariance function with damping and cutoff
kernel = 'SE_damped_cut'

# training
gp = GP_sparse_linear(bond_dict=bond_dict,
                      training_file=training_file,
                      kernel=kernel,
                      sigma=sigma,
                      rel_sigma=False)


# load literature potential 3ob for comparison
bond_dict = {'Hf-Hf': {'skf': 'fit_skf/Hf-Hf.skf'},
             'Hf-O': {'skf': 'fit_skf/Hf-O.skf'},
             'O-O': {'skf': 'fit_skf/O-O.skf'}}

tob = Spline(bond_dict)

# plotting
mod_gp = {'type': 'GP', 'model': gp,
          'label': 'GPrep', 'color': 'MediumSeaGreen'}
mod_tob = {'type': 'Spline', 'model': tob, 'label': 'CCS', 'color': 'Orange'}

mod_dict = {'gp': mod_gp, 'mio': mod_tob}

X_HfHf = np.linspace(start_HfHf, cutoff_HfHf+0.2, 200)
plot_repulsion(X_HfHf, 'Hf-Hf', mod_dict=mod_dict, der=0)

X_HfO = np.linspace(start_HfO, cutoff_HfO+0.1, 200)
plot_repulsion(X_HfO, 'Hf-O', mod_dict=mod_dict, der=0)

X_OO = np.linspace(start_OO, cutoff_OO+0.2, 200)
plot_repulsion(X_OO, 'O-O', mod_dict=mod_dict, der=0)

# spline and write to .skf
gp.make_skf_spline(s0=start_HfHf, s1=start_HfHf+0.1, N=100,
                   outputfile='Hf-Hf_GPrep_SI.skf', bond_type='Hf-Hf')
gp.make_skf_spline(s0=start_HfO, s1=start_HfO+0.1, N=100,
                   outputfile='Hf-O_GPrep_SI.skf', bond_type='Hf-O')
gp.make_skf_spline(s0=start_OO, s1=start_OO+0.1, N=100,
                   outputfile='O-O_GPrep_SI.skf', bond_type='O-O')
